import json

from langgraph.constants import END, START
from langgraph.graph import StateGraph

import os
from typing_extensions import TypedDict, Literal

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from tools import *
from state import MultiAgentState

chat_llm = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get("DEEPSEEK_API_KEY"),
                   base_url=os.environ.get("DEEPSEEK_BASE_URL"))

coder_llm = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get("DEEPSEEK_API_KEY"),
                   base_url=os.environ.get("DEEPSEEK_BASE_URL"))

db_llm = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get("DEEPSEEK_API_KEY"),
                   base_url=os.environ.get("DEEPSEEK_BASE_URL"))

# chat_llm = ChatOpenAI(model="glm-4-flash",
#                    api_key="0d28f030249b4fe38dc501510748b595.9SGm9tuJBlcgqKBm",
#                    base_url="https://open.bigmodel.cn/api/paas/v4/")
# coder_llm = ChatOpenAI(model="glm-4-flash",
#                    api_key="0d28f030249b4fe38dc501510748b595.9SGm9tuJBlcgqKBm",
#                    base_url="https://open.bigmodel.cn/api/paas/v4/")
# db_llm = ChatOpenAI(model="glm-4-flash",
#                    api_key="0d28f030249b4fe38dc501510748b595.9SGm9tuJBlcgqKBm",
#                    base_url="https://open.bigmodel.cn/api/paas/v4/")

dbTools=[add_sale,delete_sale,update_sale,query_sales]
coderTools = [python_repl]

db_agent= create_react_agent(db_llm,tools=dbTools,
    state_modifier="You use to perform database operations while should provide accurate data for the code_generator to use")
coder_agent= create_react_agent(coder_llm,tools=coderTools,
    state_modifier="Run python code to display diagrams or output execution results")

# 任何一个代理都可以决定结束
members = ["chatbot", "coder", "sqler"]
options = members + ["FINISH"]

class Router(TypedDict):
    """Worker to route to next. If no workers needed, route to FINISH"""
    next : Literal["chatbot", "coder", "sqler", "FINISH"]

def chat_node(state:MultiAgentState):
    last_message = state["messages"][-1]
    response = chat_llm.invoke(last_message.content)

    return {
        "messages": [HumanMessage(content=response.content,name="chatbot")]
    }

def db_node(state:MultiAgentState):
    response = db_agent.invoke(state)
    return {"messages":[HumanMessage(content=response["messages"][-1].content,name="sqler")]}

def coder_node(state:MultiAgentState):
    response = coder_agent.invoke(state)
    return {"messages":[HumanMessage(content=response["messages"][-1].content,name="coder")]}

def supervisor(state:MultiAgentState):
    system_prompt =  """You are a supervisor tasked with managing a conversation between the
         following workers: chatbot, coder, sqler
        Each worker has a specific role:
        - chatbot: Responds directly to user inputs using natural language.
        - coder: un python code to display diagrams or output execution results.
        - sqler: perform database operations while should provide accurate data for the code_generator to use.
         Given the following user request, respond with the worker to act next.
         Each worker will perform a task and response with their results and status.
        When you think the result has answered the user's question, just reply FINISH."""
    messages = [SystemMessage(content= system_prompt)] + state["messages"]
    response = db_llm.invoke(messages)

    try:
        response_content = response.content
        response_dict = json.loads(response_content)
        next_ = response_dict.get("next", "FINISH")
    except (json.JSONDecodeError, AttributeError):
        next_ = "FINISH"

    if next_ == "FINISH":
        next_ = END
    return {"next": next_}

workflow = StateGraph(MultiAgentState)

workflow.add_node("chatbot", chat_node)
workflow.add_node("coder", coder_node)
workflow.add_node("sqler", db_node)
workflow.add_node("supervisor", supervisor)

workflow.add_edge("chatbot","supervisor")
workflow.add_edge("coder","supervisor")
workflow.add_edge("sqler","supervisor")

workflow.add_edge(START,"supervisor")
workflow.add_conditional_edges("supervisor", lambda state: state["next"])

graph = workflow.compile()
graph.name= "multi-agent"

while True:
    input = input("User: ")
    if input == "退出":
        break
    result = graph.invoke({
        "messages": [
            HumanMessage(content=input)
        ]
    })
    print(result)
